# cython: profile=True
from numpy import array,empty,isnan,dot
from math import sqrt,pi
try:
    from mpi4py import MPI
except ImportError:
    print "cannot import MPI, continuing.."
from v3dsolver.fluid_elements.fluid_elems import PseudoParticle,Saffman,FluidElem
from v3dsolver.vector import Vec
from v3dsolver.tree_manager.spherical_t_design import get_pseudo_pos_list
from scipy.special import legendre
import pdb,time

class Cell:
    """
    Generic 3D cell class.
    """
    def __init__(self,level,cell_length,cell_center,npmax,rad,maxlevel,
                 parent=None,parentelems=[],npseudo=None):
        """
        Initialization of Cell object. Various parameters defining the cell 
        object described below:
        
        level: level of this cell on tree hierarchy
        cell_length: length of cell
        cell_center: center of cell
        npmax: max number of particles for cell to be childless
        rad: radius of blobs
        maxlevel: maximum allowable depth
        parent: parent cell of current cell (None in case of root cell)
        parentelems: elements in parent cell (All elements in case of root)
        npseudo: no. of pseudo particles in case of MakinoCell (leave default 
                 value None otherwise)
        """
        self.level = level
        if isinstance(cell_center,Vec):
            self.center = cell_center
        else:
            self.center = Vec(*cell_center)
        self.length = cell_length
        self.radius = sqrt(3)*self.length/2.0
        self.npmax = npmax
        self.rad = rad
        self.maxlevel = maxlevel
        self.parent = parent
        if parentelems==[]:
            parentelems = parent.elems
        self.parentelems = parentelems
        self.elems = self.get_elems(-1) # all elements inside
        self.cause_elems = self.get_elems(0) + self.get_elems(2)
        self.effect_elems = self.get_elems(1) + self.get_elems(2)
        self.npseudo = npseudo
        self.children = self.get_children()
        self.childless = not self.splittable()
    
    def set_id(self,id):
        """
        Set id of current cell
        """
        self.id = id
    
    def get_elems(self,cause_or_effect=-1):
        """
        Returns elements inside the cell of the type defined by the parameter 
        cause_or_effect:
        
        0: returns pure cause elements
        1: returns pure effect elements
        2: returns elements which are both cause and effect
        -1 (default): returns all elements inside the cell 
        """
        
        elems = []
        if cause_or_effect==-1: # generate self.elems if it is empty
            for elem in self.parentelems:
                if self.contains(elem):
                    elems.append(elem)
        else:
            for elem in self.elems:
                if elem.cause_or_effect==cause_or_effect:
                    elems.append(elem)
        return elems
    
    def contains(self,elem):
        """
        Returns True or False depending on whether or not elem lies inside the 
        cell. If an elem lies on the boundary, it is considered to be an elem 
        of the cell which picks it first.
        """
        elem_pos = elem.pos
        self_center = self.center
        if (abs(elem_pos.x - self_center.x)<=0.5*self.length and
            abs(elem_pos.y - self_center.y)<=0.5*self.length and
            abs(elem_pos.z - self_center.z)<=0.5*self.length):
            # elem parent for current level not decided yet
            if len(elem.parentid) is self.level: 
                # make this cell parent for current level
                elem.parentid.append(self)
            if elem.parentid[self.level] is self:
                return True
        else:
            return False

    def get_corners(self):
        """
        Returns the corners of a cell
        """
        directions = [[1.,1.,1.],[1.,1.,-1.],[1.,-1.,1.],[1.,-1.,-1.],
                      [-1.,1.,1.],[-1.,1.,-1.],[-1.,-1.,1.],[-1.,-1.,-1.]]
        corners = []
        for direction in directions:
            direction = array(direction)
            corner = self.center + self.length/2.0*direction
            corners += [corner]
        return corners

    def get_children(self):
        """
        Returns a list of new cells which are children of the current cell.
        """
        children = []
        if not self.splittable():
            pass # skip splitting for childless cells
        else:
            # directions along which the child's center could move in an 
            # oct-tree configuration
            directions = [Vec(1.,1.,1.),Vec(1.,1.,-1.),Vec(1.,-1.,1.),Vec(1.,-1.,-1.),
                         Vec(-1.,1.,1.),Vec(-1.,1.,-1.),Vec(-1.,-1.,1.),Vec(-1.,-1.,-1.)]
            for direction in directions:
                childcenter = self.center+direction*self.length/4.0
                child = self.__class__(self.level+1,self.length/2.0,
                                       childcenter,self.npmax,self.rad,
                                       self.maxlevel,parent=self,
                                       npseudo=self.npseudo)
                children += [child]
        return children
    
    def splittable(self):
        """
        Returns True if the cell has more number of cause elements than npmax, 
        otherwise False. 
        """
        if len(self.cause_elems)<=self.npmax or self.level>=self.maxlevel:
            return False
        else:
            return True
    
    def direct_vel(self,effect_elem,cause_elems=None):
        """
        Returns velocity on effect_elem calculated by this cell, directly 
        (traversing through all cause elements). If cause_elems is None, it 
        will by default calculate velocity due to all cause elements of the 
        cell, otherwise, velocity due to cause_elems supplied is returned.
        
        effect_elem: any FluidElem (or derived class) instance
        """
        vel = Vec()
        if cause_elems is None:
            cause_elems = self.cause_elems
        for cause_elem in cause_elems:
            v = cause_elem.get_vel(effect_elem)
            vel.add_inplace(v)
        return vel
    
    def is_neighbour(self,cellcenter,level):
        """
        Check whether cell defined by cellcenter and at level 'level' is nearby 
        to current cell object or not.
        """
        # length of cell defined by cellcenter
        otherlen = self.length*2**(self.level-level) 
        otherrad = sqrt(3)*otherlen*0.5
        dist = self.center - cellcenter
        distance = dist.length()
        
        if (abs(distance-otherrad)>2*self.radius):
            return False
        else:
            return True
    
    def is_well_separated(self,elem):
        """
        Checks whether the cell is well separated from a particle located at 
        elem_pos.
        """
        dist = self.center - elem.pos
        distance = dist.length()
        if distance<=2*self.radius:
            return False
        else:
            return True

class MakinoCell(Cell):
    """
    Cell based on Makino's method of pseudo particles for tree-codes in 3D 
    problems. Derived from the generic Cell class. Incorporates additional 
    methods for obtaining pseudo particles and velocity calculation based 
    on them.
    """
    def __init__(self,level,cell_length,cell_center,npmax,rad,maxlevel,
                 P_l=None,parent=None,parentelems=[],npseudo=None):
        """
        Initialization of Makino cell. Parameters except those defined below 
        same as base class Cell.
        
        npseudo: number of pseudo particles to generate 
                (if None(default), gets npseudo of parent)
        P_l: list of legendre functions till maximum limit
        # Caution: P_l and parent both cannot be None simultaneously 
        """
        if P_l==None:
            P_l = parent.P_l
        self.P_l = P_l
        Cell.__init__(self,level,cell_length,cell_center,npmax,rad,maxlevel,
                      parent,parentelems,npseudo)
        self.pseudo_p = []
        
    def get_pseudo(self,constituting_particles=[]):
        """
        Returns the pseudo particles for a cell
        
        constituting_particles: particles which contribute in formation of 
                                pseudo particles.
                                if [], cause_elems of cell form pseudoparticles
                                else, this particles list forms pseudoparticles
        """
        if constituting_particles==[]:
            constituting_particles = self.cause_elems
        if len(constituting_particles)<=self.npseudo:
            pseudo_p = constituting_particles # no point creating pseudo particles
        else:
            pseudo_p = []
            pseudo_pos_list = get_pseudo_pos_list(self,self.npseudo)
            for pos in pseudo_pos_list:
                newpseudo = PseudoParticle(pos,parent_cell=self,rad=self.rad,
                            P_l=self.P_l,particles=constituting_particles) 
                pseudo_p += [newpseudo]
        if constituting_particles==self.cause_elems:
            self.pseudo_p = pseudo_p
        return pseudo_p
    
    def transfer_pseudo(self):
        """
        Returns the pseudo particles for this cell from it's children's pseudo 
        particles.
        """
        pseudo_elems = []
        for child in self.children:
            pseudo_elems += child.pseudo_p
        # creating pseudo particles of self using it's children's pseudo 
        # particles
        self.pseudo_p = self.get_pseudo(pseudo_elems)
    
    def compute_vel(self,elem):
        """
        Returns velocity at elem by traversing tree starting at current 
        cell.
        
        elem: any effect element
        """
        vel = Vec()
        if self.is_well_separated(elem):
            vel = self.direct_vel(elem,self.pseudo_p)
        else:
            if self.childless:
                vel = self.direct_vel(elem)
            else:
                for child in self.children:
                    v = child.compute_vel(elem)
                    vel.add_inplace(v)
        return vel

class TreeManager:
    """
    Generic tree manager class which distributes the elements among different 
    processors, builds the tree and performs velocity computations, in a 
    parallel fashion.
    """ 
    def __init__(self,elems,domain,npmax,npseudo,rad,legendre_lim=4,maxlevel=7,comm=None):
        """
        Initialization of a TreeManager object.
        
        comm: communicator object for parallel computation (Keep default None
              for serial computation)
        elems: list of elements of current processor
        domain: global domain of the form [xmin,ymin,zmin,length] based on 
                elements of all processors
        npmax: max number of non-effect elements a childless cell can have
        npseudo: number of pseudo particles to introduce in each cell
        legendre_lim: maximum degree of legendre polynomials to be used in pseudo
                      particle generation
        rad: radius of blobs 
        """
        self.comm = comm
        self.npmax = npmax
        self.domain = domain # global domain
        self.elems = elems
        rootcenter = array(self.domain[:3])+0.5*self.domain[3]*array([1,1,1])
        self.npseudo = npseudo
        self.maxlevel = maxlevel
        self.P_l = []
        for l in range(legendre_lim):
            self.P_l += [list(legendre(l).coeffs[::-1])]
        self.root = MakinoCell(level=0,cell_length=self.domain[3],
                               cell_center=rootcenter,npmax=self.npmax,
                               rad=rad,P_l=self.P_l,parent=None,parentelems=self.elems,
                               npseudo=self.npseudo,maxlevel=maxlevel)
        self.childless_cells = []
        self.rad = rad
    
    def build_tree(self):
        """
        Builds tree by recursively splitting cells if they aren't childless and
        adding childless cells to the tree manager record. Also sets the 
        maxlevel of tree manager object in the attribute maxlevel.
        """
        self.cells = [[self.root]]
        level = 0
        while not len(self.cells[level])==0:
            self.cells += [[]]
            for cell in self.cells[level]:
                if not cell.childless:
                    self.cells[level+1] += cell.children
                else: # childless
                    self.childless_cells += [cell]
            level += 1
        self.maxlevel = level - 1 # max level of tree
        self.set_cell_ids()
        for elem in self.elems:
            elem.parentid = []
        
    def set_cell_ids(self):
        """
        Sets ids for all the cells in tree
        """
        id = 0
        for level in range(len(self.cells)):
            for cell in self.cells[level]:
                cell.set_id(id)
                id += 1
    
    def get_cell(self,id):
        """
        Returns cell object corresponding to the specified id
        """
        pointer = 0
        level = 0
        while self.cells[level][0].id <=id:
            level += 1
            if level==self.maxlevel+1:
                break
        # level has now crossed id
        index = id - self.cells[level-1][0].id
        return self.cells[level-1][index]
    
    def generate_pseudo_particles(self):
        """
        Obtains pseudo particles for cell at all levels.
        """
        for cell in self.childless_cells:
            cell.get_pseudo()
        for level in range(self.maxlevel-1,-1,-1):
            for cell in self.cells[level]:
                cell.transfer_pseudo() # obtain pseudo particles from children
